import pywt
from scipy.signal import butter, filtfilt

def wavelet_transform(trace, threshold=0.2):
        w = pywt.Wavelet("db8")
        max_level = pywt.dwt_max_level(len(trace), w.dec_len)
        coeffs = pywt.wavedec(trace, "db8", level=max_level)
        for i in range(1, len(coeffs)):
            coeffs[i] = pywt.threshold(coeffs[i], threshold * max(coeffs[i]))
        trace_reconfig = pywt.waverec(coeffs, "db8")
        return trace_reconfig

def lowpass_filter(traces,order, ncf):
        b, a = butter(order, ncf, btype='low', analog=False, output='ba')
        return filtfilt(b, a, traces)

def bandpass_filter(traces,order, ncf):
        b, a = butter(order, ncf, btype='band', analog=False, output='ba')
        return filtfilt(b, a, traces)